{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In CoreML Neural Network Specification version 4 (which is available from iOS 13 and MacOS 10.15), several \"control-flow\" layers have been added. CoreML spec is described in the protobuf format and for a list of all supported layer types and documentation, see [here](https://github.com/apple/coremltools/blob/master/mlmodel/format/NeuralNetwork.proto).\n",
    "\n",
    "In this notebook, we build a neural network that uses a few of the new control flow layers. We will write a simple python program to compute the largest eigenvalue of a given matrix and then show how a neural network can be built to replicate that program in an mlmodel.\n",
    "\n",
    "We choose the [power iteration method](https://en.wikipedia.org/wiki/Power_iteration). It is a simple iterative algorithm. Given a square matrix, $A$ of dimensions $n\\times n$, it computes the largest eigenvalue (by magnitude) and the corresponding eigenvector (the algorithm can be adapted to compute all the eigenvalues, however we do not implement that here). \n",
    "\n",
    "Here is how the algorithm works. Pick a normalized random vector to start with, $x$, of dimension $n$. Repetitively, multiply it by the matrix and normalize it, i.e., $x\\leftarrow Ax$ and $x\\leftarrow \\frac{x}{\\left \\| x \\right \\|}$. Gradually the vector converges to the largest eigenvector. Simple as that! \n",
    "There are a few conditions that the matrix should satisfy for this to happen, but let us not worry about it for this example. \n",
    "For now we will assume that the matrix is real and symmetric, this guarantees the eigenvalues to be real. \n",
    "After we have the normalized eigenvector, the corresponding eigenvalue can be computed by the formula $x^TAx$\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's code this up in Python using Numpy!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0: diff: 6.69187030143e-05\n",
      "1: diff: 0.00208718410489\n",
      "2: diff: 0.0614522880272\n",
      "3: diff: 0.771617699317\n",
      "4: diff: 0.193129218664\n",
      "5: diff: 0.0075077446807\n",
      "6: diff: 0.000241962094403\n",
      "7: diff: 7.74407193072e-06\n",
      "8: diff: 2.47796068775e-07\n",
      "Largest eigenvalue: 8.5249 \n",
      "('Corresponding eigenvector: ', array([-0.74152421,  0.67092611]))\n"
     ]
    }
   ],
   "source": [
    "import numpy as np\n",
    "import copy\n",
    "\n",
    "np.random.seed(8) # try different seeds to play with the number of iterations it takes for convergence!\n",
    "\n",
    "'''\n",
    "Use power method to compute the largest eigenvalue of a real symmetric matrix\n",
    "'''\n",
    "\n",
    "convergence_tolerance = 1e-6 # decrease/increase to trade off precision\n",
    "number_of_iterations = 100 # decrease/increase to trade off precision\n",
    "\n",
    "def power_iteration(matrix, starting_vector):\n",
    "    x = copy.deepcopy(starting_vector)\n",
    "    for i in range(number_of_iterations):\n",
    "        y = np.matmul(A,x)\n",
    "        #normalize\n",
    "        y = y / np.sqrt(np.sum(y**2))\n",
    "        # compute the diff to check for convergence\n",
    "        # we use cosine difference as both vectors are normalized and can get\n",
    "        # rotated by 180 degrees between iterations\n",
    "        diff = 1-abs(np.dot(x,y))\n",
    "        # update x\n",
    "        x = y\n",
    "        print('{}: diff: {}'.format(i, diff))\n",
    "        if diff < convergence_tolerance: \n",
    "            break\n",
    "\n",
    "    x_t = np.transpose(x)\n",
    "    eigen_value = np.matmul(x_t, np.matmul(A,x))\n",
    "    return eigen_value, x\n",
    "    \n",
    "\n",
    "# define the symmetric real matrix for which we need the eigenvalue. \n",
    "A = np.array([[4,-5], [-5,3]], dtype=np.float)\n",
    "\n",
    "# a random starting vector\n",
    "starting_vector = np.random.rand(2)\n",
    "starting_vector = starting_vector / np.sqrt(np.sum(starting_vector**2)) ## normalize it\n",
    " \n",
    "eigen_value, eigen_vector = power_iteration(A, starting_vector)\n",
    "\n",
    "print('Largest eigenvalue: %.4f ' % eigen_value)\n",
    "print('Corresponding eigenvector: ', eigen_vector)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We see that in this case, the algorithm converged, given our specified toelrance, in 9 iterations. \n",
    "To confirm whether the eigenvalue is correct, lets use the \"linalg\" sub-package of numpy. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "numpy linalg: largest eigenvalue: 8.5249 \n",
      "('numpy linalg: first eigenvector: ', array([ 0.74145253, -0.67100532]))\n"
     ]
    }
   ],
   "source": [
    "from numpy import linalg as LA\n",
    "\n",
    "e, v = LA.eig(A)\n",
    "idx = np.argmax(abs(e))\n",
    "print('numpy linalg: largest eigenvalue: %.4f ' % e[idx])\n",
    "print('numpy linalg: first eigenvector: ', v[:,idx])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Indeed we see that the eigenvalue matches with our power iteration code. The eigenvector is rotated by 180 degrees, but that is fine.\n",
    "\n",
    "Now, lets build an mlmodel to do the same. We use the builder API provided by coremltools to write out the protobuf messages. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "import coremltools\n",
    "import coremltools.models.datatypes as datatypes\n",
    "from coremltools.models.neural_network import NeuralNetworkBuilder\n",
    "\n",
    "input_features = [('matrix', datatypes.Array(*(2,2))),\n",
    "                  ('starting_vector', datatypes.Array(*(2,)))]\n",
    "\n",
    "output_features = [('maximum_eigen_value', datatypes.Array(*(1,))), \n",
    "                   ('eigen_vector', None),\n",
    "                   ('iteration_count', datatypes.Array(*(1,)))]\n",
    "\n",
    "builder = NeuralNetworkBuilder(input_features, output_features, disable_rank5_shape_mapping=True)\n",
    "\n",
    "# convert the starting_vector which has shape (2,) to shape (2,1) \n",
    "# so that it can be used by the Batched-MatMul layer\n",
    "builder.add_expand_dims('expand_dims', 'starting_vector', 'x', axes=[-1])\n",
    "builder.add_load_constant_nd('iteration_count', 'iteration_count',\n",
    "                             constant_value=np.zeros((1,)),\n",
    "                             shape=(1,))\n",
    "\n",
    "# start building the loop\n",
    "loop_layer = builder.add_loop('loop', max_iterations=number_of_iterations)\n",
    "# get the builder object for the \"body\" of the loop\n",
    "loop_body_builder = NeuralNetworkBuilder(nn_spec=loop_layer.loop.bodyNetwork)\n",
    "\n",
    "# matrix multiply\n",
    "# input shapes: (n,n),(n,1)\n",
    "# output shape: (n,1)\n",
    "loop_body_builder.add_batched_mat_mul('bmm.1', input_names=['matrix','x'], output_name='y')\n",
    "# normalize the vector\n",
    "loop_body_builder.add_reduce_l2('reduce', input_name='y', output_name='norm', axes = [0])\n",
    "loop_body_builder.add_divide_broadcastable('divide', ['y','norm'], 'y_normalized')\n",
    "\n",
    "# find difference with previous, which is computed as (1 - abs(cosine diff))\n",
    "loop_body_builder.add_batched_mat_mul('cosine', ['y_normalized', 'x'], 'cosine_diff', transpose_a=True)\n",
    "loop_body_builder.add_unary('abs_cosine','cosine_diff','abs_cosine_diff', mode='abs')\n",
    "loop_body_builder.add_activation('diff', non_linearity='LINEAR',\n",
    "                                 input_name='abs_cosine_diff',\n",
    "                                 output_name='diff', params=[-1,1])\n",
    "\n",
    "# update iteration count\n",
    "loop_body_builder.add_activation('iteration_count_add', non_linearity='LINEAR',\n",
    "                                 input_name='iteration_count',\n",
    "                                 output_name='iteration_count_plus_1', params=[1,1])\n",
    "loop_body_builder.add_copy('iteration_count_update', 'iteration_count_plus_1', 'iteration_count')\n",
    "\n",
    "# update 'x'\n",
    "loop_body_builder.add_copy('update_x', 'y_normalized', 'x')\n",
    "\n",
    "# add condition to break from the loop, if convergence criterion is met\n",
    "loop_body_builder.add_less_than('cond', ['diff'], 'cond', alpha=convergence_tolerance)\n",
    "branch_layer = loop_body_builder.add_branch('branch_layer', 'cond')\n",
    "builder_ifbranch = NeuralNetworkBuilder(nn_spec=branch_layer.branch.ifBranch)\n",
    "builder_ifbranch.add_loop_break('break')\n",
    "\n",
    "# now we are out of the loop, compute the eigenvalue\n",
    "builder.add_batched_mat_mul('bmm.2', input_names=['matrix','x'], output_name='x_right')\n",
    "builder.add_batched_mat_mul('bmm.3', input_names=['x','x_right'], output_name='maximum_eigen_value', transpose_a=True)\n",
    "builder.add_squeeze('squeeze', 'x', 'eigen_vector', squeeze_all=True)\n",
    "\n",
    "spec = builder.spec\n",
    "model = coremltools.models.MLModel(spec)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Okay, so now we have the mlmodel spec. Before we call predict on it, lets print it out to check whether everything looks okay. We use the utility called \"print_network_spec\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Inputs:\n",
      "  matrix [2, 2]\n",
      "  starting_vector [2]\n",
      "Outputs:\n",
      "  maximum_eigen_value [1]\n",
      "  eigen_vector []\n",
      "  iteration_count [1]\n",
      "\n",
      "\n",
      "def model(matrix, starting_vector) :\n",
      "\tx = \u001b[91m expandDims\u001b[00m\u001b[94m (starting_vector)\u001b[00m\n",
      "\titeration_count = \u001b[91m loadConstantND\u001b[00m\u001b[94m (shape = \u001b[00m(1,), \u001b[94m value = \u001b[00m[0.0]\u001b[94m )\u001b[00m\n",
      "\u001b[91m \tloop\u001b[00m\u001b[94m ()\u001b[00m\n",
      "\t\ty = \u001b[91m batchedMatmul\u001b[00m\u001b[94m (matrix, x)\u001b[00m\n",
      "\t\tnorm = \u001b[91m reduceL2\u001b[00m\u001b[94m (y)\u001b[00m\n",
      "\t\ty_normalized = \u001b[91m divideBroadcastable\u001b[00m\u001b[94m (y, norm)\u001b[00m\n",
      "\t\tcosine_diff = \u001b[91m batchedMatmul\u001b[00m\u001b[94m (y_normalized, x)\u001b[00m\n",
      "\t\tabs_cosine_diff = \u001b[91m unary\u001b[00m\u001b[94m (cosine_diff)\u001b[00m\n",
      "\t\tdiff = \u001b[91m activation\u001b[00m\u001b[94m (abs_cosine_diff)\u001b[00m\n",
      "\t\titeration_count_plus_1 = \u001b[91m activation\u001b[00m\u001b[94m (iteration_count)\u001b[00m\n",
      "\t\titeration_count = \u001b[91m copy\u001b[00m\u001b[94m (iteration_count_plus_1)\u001b[00m\n",
      "\t\tx = \u001b[91m copy\u001b[00m\u001b[94m (y_normalized)\u001b[00m\n",
      "\t\tcond = \u001b[91m lessThan\u001b[00m\u001b[94m (diff)\u001b[00m\n",
      "\u001b[91m \t\tbranch\u001b[00m\u001b[94m (cond)\u001b[00m\n",
      "\u001b[91m \t\tIfBranch:\u001b[00m\n",
      "\u001b[91m \t\t\tloopBreak\u001b[00m\n",
      "\tx_right = \u001b[91m batchedMatmul\u001b[00m\u001b[94m (matrix, x)\u001b[00m\n",
      "\tmaximum_eigen_value = \u001b[91m batchedMatmul\u001b[00m\u001b[94m (x, x_right)\u001b[00m\n",
      "\teigen_vector = \u001b[91m squeeze\u001b[00m\u001b[94m (x)\u001b[00m\n",
      "\u001b[91m \n",
      "\treturn \u001b[00mmaximum_eigen_value, eigen_vector, iteration_count\n"
     ]
    }
   ],
   "source": [
    "from  coremltools.models.neural_network.printer import print_network_spec\n",
    "print_network_spec(spec, style='coding')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CoreML computed eigenvalue: 8.5249\n",
      "('CoreML computed eigenvector: ', array([-0.74152416,  0.67092603]), (2,))\n",
      "CoreML iteration count: 9\n"
     ]
    }
   ],
   "source": [
    "# call predict on CoreML model\n",
    "input_dict = {}\n",
    "input_dict['starting_vector'] = starting_vector\n",
    "input_dict['matrix'] = A.astype(np.float)\n",
    "\n",
    "output = model.predict(input_dict)\n",
    "coreml_eigen_value = output['maximum_eigen_value']\n",
    "coreml_eigen_vector = output['eigen_vector']\n",
    "\n",
    "print('CoreML computed eigenvalue: %.4f' % coreml_eigen_value)\n",
    "print('CoreML computed eigenvector: ', coreml_eigen_vector, coreml_eigen_vector.shape)\n",
    "print('CoreML iteration count: %d' % output['iteration_count'])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Indeed the output matches with our python program. \n",
    "\n",
    "Although, we do not do it here, the parameters \"convergence_tolerance\" and \"number_of_iterations\" can be made as network inputs, so that their value can be modifed at runtime. \n",
    "\n",
    "Currently, the input shapes to the Core ML model are fixed, $(2, 2)$ for the matrix and $(2,)$ for the starting vector. However, we can add shape flexibility so that the same mlmodel can be run on different input sizes. There are two ways to specify shape flexibility, either through \"ranges\" or via a list of \"enumerated\" shapes. Here we specify the latter."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "from coremltools.models.neural_network import flexible_shape_utils\n",
    "\n",
    "# (2,2) has already been provided as the default shape for \"matrix\" \n",
    "# during initialization of the builder,\n",
    "# here we add two more shapes that will be allowed at runtime\n",
    "flexible_shape_utils.add_multiarray_ndshape_enumeration(spec, \n",
    "                                                        feature_name='matrix',\n",
    "                                                        enumerated_shapes=[(3,3), (4,4)])\n",
    "\n",
    "# (2,) has already been provided as the default shape for \"matrix\" \n",
    "# during initialization of the builder,\n",
    "# here we add two more shapes that will be allowed at runtime\n",
    "flexible_shape_utils.add_multiarray_ndshape_enumeration(spec, \n",
    "                                                        feature_name='starting_vector',\n",
    "                                                        enumerated_shapes=[(3,), (4,)])\n",
    "\n",
    "model = coremltools.models.MLModel(spec)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0: diff: 0.99757552989\n",
      "1: diff: 0.718149467089\n",
      "2: diff: 0.492558374678\n",
      "3: diff: 0.325410135011\n",
      "4: diff: 0.208606358183\n",
      "5: diff: 0.130795340624\n",
      "6: diff: 0.0807677916817\n",
      "7: diff: 0.0493798553633\n",
      "8: diff: 0.0299993308647\n",
      "9: diff: 0.0181536364413\n",
      "10: diff: 0.0109588786353\n",
      "11: diff: 0.00660585926588\n",
      "12: diff: 0.0039783687005\n",
      "13: diff: 0.00239467498795\n",
      "14: diff: 0.00144094325621\n",
      "15: diff: 0.000866886171118\n",
      "16: diff: 0.000521466038849\n",
      "17: diff: 0.00031366000502\n",
      "18: diff: 0.000188657339187\n",
      "19: diff: 0.000113468967192\n",
      "20: diff: 6.82454629412e-05\n",
      "21: diff: 4.1045582895e-05\n",
      "22: diff: 2.46863363353e-05\n",
      "23: diff: 1.48472285797e-05\n",
      "24: diff: 8.92962598664e-06\n",
      "25: diff: 5.37057288463e-06\n",
      "26: diff: 3.23003808245e-06\n",
      "27: diff: 1.94264962894e-06\n",
      "28: diff: 1.16837216313e-06\n",
      "29: diff: 7.02696602684e-07\n",
      "python code: largest eigenvalue: -11.7530 \n",
      "('python code: corresponding eigenvector: ', array([ 0.61622756,  0.52125649, -0.59038569]))\n"
     ]
    }
   ],
   "source": [
    "# lets run the model with a (3,3) matrix \n",
    "A = np.array([[1, -6, 8], [-6, 1, 5], [8, 5, 1]], dtype=np.float)\n",
    "\n",
    "starting_vector = np.random.rand(3)\n",
    "starting_vector = starting_vector / np.sqrt(np.sum(starting_vector**2)) ## normalize it\n",
    "\n",
    "eigen_value, eigen_vector = power_iteration(A, starting_vector)\n",
    "\n",
    "print('python code: largest eigenvalue: %.4f ' % eigen_value)\n",
    "print('python code: corresponding eigenvector: ', eigen_vector)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "numpy linalg: largest eigenvalue: -11.7530 \n",
      "('numpy linalg: first eigenvector: ', array([-0.61583909, -0.5213392 ,  0.59071791]))\n"
     ]
    }
   ],
   "source": [
    "from numpy import linalg as LA\n",
    "\n",
    "e, v = LA.eig(A)\n",
    "idx = np.argmax(abs(e))\n",
    "print('numpy linalg: largest eigenvalue: %.4f ' % e[idx])\n",
    "print('numpy linalg: first eigenvector: ', v[:,idx])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CoreML computed eigenvalue: -11.7530\n",
      "('CoreML computed eigenvector: ', array([ 0.61622757,  0.52125645, -0.59038568]), (3,))\n",
      "CoreML iteration count: 30\n"
     ]
    }
   ],
   "source": [
    "input_dict['starting_vector'] = starting_vector\n",
    "input_dict['matrix'] = A.astype(np.float)\n",
    "\n",
    "output = model.predict(input_dict)\n",
    "coreml_eigen_value = output['maximum_eigen_value']\n",
    "coreml_eigen_vector = output['eigen_vector']\n",
    "\n",
    "print('CoreML computed eigenvalue: %.4f' % coreml_eigen_value)\n",
    "print('CoreML computed eigenvector: ', coreml_eigen_vector, coreml_eigen_vector.shape)\n",
    "print('CoreML iteration count: %d' % output['iteration_count'])"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
